import os
import torch
import numpy as np
import random
from glob import glob
from torch.utils.data import Dataset
from torchvision import transforms
import h5py
import itertools
from torch.utils.data.sampler import Sampler
import torch.nn.functional as F

class LAHeart(Dataset):
    """ LA Dataset """
    def __init__(self, base_dir=None, split='train', num=None, transform=None,strong_transform=None):
        self._base_dir = base_dir
        self.transform = transform
        self.strong_transform = strong_transform
        self.sample_list = []
        if split=='train':
            with open(self._base_dir+'/../train.list', 'r') as f:
                self.image_list = f.readlines()
        elif split == 'test':
            with open(self._base_dir+'/../test.list', 'r') as f:
                self.image_list = f.readlines()
        self.image_list = [item.replace('\n','') for item in self.image_list]
        if num is not None:
            self.image_list = self.image_list[:num]
        print("total {} samples".format(len(self.image_list)))

    def __len__(self):
        return len(self.image_list)

    def __getitem__(self, idx):
        image_name = self.image_list[idx]
        h5f = h5py.File(self._base_dir+"/"+image_name+"/mri_norm2.h5", 'r')
        image = h5f['image'][:]
        label = h5f['label'][:]
        sample = {'image': image, 'label': label}
        if self.transform: #弱图像增强
            sample = self.transform(sample)
            weak_image, weak_label = sample['image'], sample['label']

        if self.strong_transform:#强图像增强
            sample = self.strong_transform(sample)

            # kernel_size = int(random.random() * 4.95)
            # kernel_size = kernel_size + 1 if kernel_size % 2 == 0 else kernel_size
            # blurring_image = transforms.GaussianBlur(kernel_size, sigma=(0.1, 2.0))#高斯模糊
            # color_jitter = transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.25)#颜色抖动

            # if random.random() < 0.8: #80%概率应用颜色抖动,标签应该就不用颜色抖动了吧？
            #     sample['image']= color_jitter(sample['image'])
            #     sample['image']= transforms.RandomGrayscale(p=0.2)(sample['image'])

            # if random.random() < 0.5: #50%概率应用高斯模糊
            #     sample['image'] = blurring_image(sample['image'])

            strong_image, strong_label = sample['image'], sample['label']

        return {
            'weak_image': weak_image,
            'weak_label': weak_label,
            'strong_image': strong_image,
            'strong_label': strong_label
        }

class CenterCrop(object):
    def __init__(self, output_size):
        self.output_size = output_size

    def __call__(self, sample):
        image, label = sample['image'], sample['label']

        # pad the sample if necessary
        if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \
                self.output_size[2]:
            pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0)
            ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0)
            pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0)
            image = np.pad(image, [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0)
            label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0)

        (w, h, d) = image.shape

        w1 = int(round((w - self.output_size[0]) / 2.))
        h1 = int(round((h - self.output_size[1]) / 2.))
        d1 = int(round((d - self.output_size[2]) / 2.))

        label = label[w1:w1 + self.output_size[0], h1:h1 + self.output_size[1], d1:d1 + self.output_size[2]]
        image = image[w1:w1 + self.output_size[0], h1:h1 + self.output_size[1], d1:d1 + self.output_size[2]]

        return {'image': image, 'label': label}


class RandomCrop(object):
    """
    Crop randomly the image in a sample
    Args:
    output_size (int): Desired output size
    """

    def __init__(self, output_size):
        self.output_size = output_size

    def __call__(self, sample):
        image, label = sample['image'], sample['label']

        # pad the sample if necessary
        if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \
                self.output_size[2]:
            pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0)
            ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0)
            pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0)
            image = np.pad(image, [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0)
            label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0)

        (w, h, d) = image.shape
        # if np.random.uniform() > 0.33:
        #     w1 = np.random.randint((w - self.output_size[0])//4, 3*(w - self.output_size[0])//4)
        #     h1 = np.random.randint((h - self.output_size[1])//4, 3*(h - self.output_size[1])//4)
        # else:
        w1 = np.random.randint(0, w - self.output_size[0])
        h1 = np.random.randint(0, h - self.output_size[1])
        d1 = np.random.randint(0, d - self.output_size[2])

        label = label[w1:w1 + self.output_size[0], h1:h1 + self.output_size[1], d1:d1 + self.output_size[2]]
        image = image[w1:w1 + self.output_size[0], h1:h1 + self.output_size[1], d1:d1 + self.output_size[2]]
        return {'image': image, 'label': label}


class RandomRotFlip(object):
    """
    Crop randomly flip the dataset in a sample
    Args:
    output_size (int): Desired output size
    """

    def __call__(self, sample):
        image, label = sample['image'], sample['label']
        k = np.random.randint(0, 4)
        image = np.rot90(image, k)
        label = np.rot90(label, k)
        axis = np.random.randint(0, 2)
        image = np.flip(image, axis=axis).copy()
        label = np.flip(label, axis=axis).copy()

        return {'image': image, 'label': label}


class RandomNoise(object):
    def __init__(self, mu=0, sigma=0.1):
        self.mu = mu
        self.sigma = sigma

    def __call__(self, sample):
        image, label = sample['image'], sample['label']
        noise = np.clip(self.sigma * np.random.randn(image.shape[0], image.shape[1], image.shape[2]), -2*self.sigma, 2*self.sigma)
        noise = noise + self.mu
        image = image + noise
        return {'image': image, 'label': label}


class ColorJitter3D:
    def __init__(self, brightness=0, contrast=0):
        """
        3D 图像的颜色抖动（仅支持亮度和对比度调整）
        :param brightness: 亮度调整范围（如 0.2 表示亮度在 80% 到 120% 之间随机调整）
        :param contrast: 对比度调整范围（如 0.2 表示对比度在 80% 到 120% 之间随机调整）
        """
        self.brightness = brightness
        self.contrast = contrast

    def __call__(self, sample):
        """
        :param img: 输入的 3D 图像，形状为 (D, H, W)
        :return: 颜色抖动后的 3D 图像
        :80%概率应用颜色抖动
        """
        image, label = sample['image'], sample['label']
        if self.brightness > 0:
            brightness_factor = torch.tensor(1.0).uniform_(1 - self.brightness, 1 + self.brightness).item()
            image = image * brightness_factor

        if self.contrast > 0 and random.random() < 0.8:
            contrast_factor = torch.tensor(1.0).uniform_(1 - self.contrast, 1 + self.contrast).item()
            mean = image.mean()
            image = (image - mean) * contrast_factor + mean

        return {'image': image, 'label': label}


class GaussianBlur3D:
    def __init__(self, kernel_size=3, sigma=(0.1, 2.0)):
        """
        3D 高斯模糊
        :param kernel_size: 高斯核大小（必须是奇数）
        :param sigma: 高斯核标准差范围（如 (0.1, 2.0) 表示标准差在 0.1 到 2.0 之间随机选择）
        """
        self.kernel_size = kernel_size
        self.sigma = sigma

    def __call__(self, sample):
        """
        :param img: 输入的 3D 图像，形状为 (D, H, W)，需要是张量类型
        :return: 高斯模糊后的 3D 图像
        :50%概率应用高斯模糊
        """
        image, label = sample['image'], sample['label']
        #print("image shape:",image.shape)[1,112,112,80]
        # 生成 3D 高斯核
        if random.random() < 0.5:
            kernel = self._get_gaussian_kernel3d(self.kernel_size, self.sigma)
            kernel = kernel.to(image.device)

            # 对 3D 图像进行卷积
            image = image.unsqueeze(0)  # 增加 batch 和 channel 维度
            image = F.conv3d(image, kernel, padding=self.kernel_size // 2)
            image = image.squeeze(0)  # 移除 batch 和 channel 维度
        return {'image': image, 'label': label}

    def _get_gaussian_kernel3d(self, kernel_size, sigma_range):
        """
        生成 3D 高斯核
        :param kernel_size: 高斯核大小
        :param sigma_range: 高斯核标准差范围
        :return: 3D 高斯核，形状为 (1, 1, kernel_size, kernel_size, kernel_size)
        """
        sigma = torch.empty(1).uniform_(sigma_range[0], sigma_range[1]).item()
        kernel_1d = torch.arange(-kernel_size // 2 + 1, kernel_size // 2 + 1, dtype=torch.float32)
        kernel_1d = torch.exp(-kernel_1d ** 2 / (2 * sigma ** 2))
        kernel_1d = kernel_1d / kernel_1d.sum()

        # 生成 3D 高斯核
        kernel_3d = torch.einsum('i,j,k->ijk', kernel_1d, kernel_1d, kernel_1d)
        kernel_3d = kernel_3d / kernel_3d.sum()
        kernel_3d = kernel_3d.unsqueeze(0).unsqueeze(0)  # 增加 channel 和 batch 维度
        return kernel_3d

class CreateOnehotLabel(object):
    def __init__(self, num_classes):
        self.num_classes = num_classes

    def __call__(self, sample):
        image, label = sample['image'], sample['label']
        onehot_label = np.zeros((self.num_classes, label.shape[0], label.shape[1], label.shape[2]), dtype=np.float32)
        for i in range(self.num_classes):
            onehot_label[i, :, :, :] = (label == i).astype(np.float32)
        return {'image': image, 'label': label,'onehot_label':onehot_label}


class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample):
        image = sample['image']
        image = image.reshape(1, image.shape[0], image.shape[1], image.shape[2]).astype(np.float32)
        if 'onehot_label' in sample:
            return {'image': torch.from_numpy(image), 'label': torch.from_numpy(sample['label']).long(),
                    'onehot_label': torch.from_numpy(sample['onehot_label']).long()}
        else:
            return {'image': torch.from_numpy(image), 'label': torch.from_numpy(sample['label']).long()}


class TwoStreamBatchSampler(Sampler):
    """Iterate two sets of indices

    An 'epoch' is one iteration through the primary indices.
    During the epoch, the secondary indices are iterated through
    as many times as needed.
    """
    def __init__(self, primary_indices, secondary_indices, batch_size, secondary_batch_size):
        self.primary_indices = primary_indices
        self.secondary_indices = secondary_indices
        self.secondary_batch_size = secondary_batch_size
        self.primary_batch_size = batch_size - secondary_batch_size

        self.current_batch_indices = []

        assert len(self.primary_indices) >= self.primary_batch_size > 0
        assert len(self.secondary_indices) >= self.secondary_batch_size > 0

    def __iter__(self):
        primary_iter = iterate_once(self.primary_indices)
        secondary_iter = iterate_eternally(self.secondary_indices)
        return (
            primary_batch + secondary_batch
            for (primary_batch, secondary_batch)
            in zip(grouper(primary_iter, self.primary_batch_size),
                    grouper(secondary_iter, self.secondary_batch_size))
        )

    def __len__(self):
        return len(self.primary_indices) // self.primary_batch_size

def iterate_once(iterable):
    return np.random.permutation(iterable)


def iterate_eternally(indices):
    def infinite_shuffles():
        while True:
            yield np.random.permutation(indices)
    return itertools.chain.from_iterable(infinite_shuffles())


def grouper(iterable, n):
    "Collect data into fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3) --> ABC DEF"
    args = [iter(iterable)] * n
    return zip(*args)
